import pickle
from tqdm import tqdm
from transformers import BertTokenizerFast


def data_preprocess(txt_path, pkl_path):
    # 读取文件
    with open(txt_path, 'r', encoding='utf-8') as f:
        data = f.read()
    # print('原始数据：',data)
    # 断指再植血管危象的手术治疗有些什么？
    # 断指再植；静脉皮瓣游离移植
    #
    # 小孩发烧又吐怎么办？，无
    # 如果持续的发烧发炎的话, 用青霉素肌注或者静脉注射都可以.效果很好的.如果发炎好了以后还是偏大, 一般得孩子长大后会变小。估计你的孩子退烧后还没有完全好吧。多喝水，吃一些口服的清热去火的药物，少吃辣的和甜的，可以减少复发

    # 加载分词器
    tokenizer = BertTokenizerFast('../vocab/vocab.txt',
                                  sep_token="[SEP]",
                                  pad_token="[PAD]",
                                  cls_token="[CLS]")
    print('分词器词表大小：', tokenizer.vocab_size)  # 13317

    cls_id = tokenizer.cls_token_id
    sep_id = tokenizer.sep_token_id

    # 数据样本分割
    if "\r\n" in data:
        train_data = data.split("\r\n\r\n")
    else:
        train_data = data.split("\n\n")

    # print('样本分割后的数据：',train_data[1:3])
    # ['断指再植血管危象的手术治疗有些什么？\n断指再植；静脉皮瓣游离移植', '小孩发烧又吐怎么办？，无\n如果持续的发烧发炎的话,用青霉素肌注或者静脉注射都可以.效果很好的.如果发炎好了以后还是偏大,一般得孩子长大后会变小。估计你的孩子退烧后还没有完全好吧。多喝水，吃一些口服的清热去火的药物，少吃辣的和甜的，可以减少复发']

    # 中文文本转换成数字文本，两句话转换成一句话。
    # [ [ [CLS]seq1[SEP]seq2[SEP] ], [ [CLS]seq1[SEP]seq2[SEP] ],...]
    chat_list = []

    # 单条数据QA分割
    for chat in tqdm(train_data):
        if "\r\n" in chat:
            sentences = chat.split("\r\n")
        else:
            sentences = chat.split("\n")
        # print(sentences) # ['断指再植血管危象的手术治疗有些什么？', '断指再植；静脉皮瓣游离移植']['Q','A']

        # 先在列表中加入一个cls_id，再遍历QA，先cls_id+Q，再append sep_id，组成cls_id+Q+sep_id，再+A+sep_id
        input_ids = [cls_id]

        # 数据数值化并加特殊token
        for sentence in sentences:
            # 数值化
            encode_seq = tokenizer.encode(sentence, add_special_tokens=False)

            # 加特殊token
            input_ids += encode_seq
            input_ids.append(sep_id)

        chat_list.append(input_ids)
    with open(pkl_path, 'wb') as f:
        pickle.dump(chat_list, f)


if __name__ == '__main__':
    txt_path = '../data/medical_train.txt'
    pkl_path = '../data/medical_train.pkl'
    data_preprocess(txt_path, pkl_path)
